import pandas as pd
import os
import numpy as np
import info
import re
import time
import matplotlib.pyplot as plt
import GetExternalDataSourceModule_CLIENT
import math

class DataLoaderModule:
    def __init__(self):
        self.connector = GetExternalDataSourceModule_CLIENT.GetExternalDataSource_CLIENT()

        self.contract_info = pd.read_pickle(
            info.DATAPATH.dataRootPath + '\\' + 'future_basic_data' + '\\' + 'ContractInfo.pkl')
        self.calendar = self.connector.get_file('tradingCalendar', info.DATAPATH.server_root_path + '/' +
                                                'calendar' + '/' + 'tradingCalendar.pkl')
        self.calendar.index = self.calendar.TradingDate.tolist()

    def read_hfData(self, product_bardata_path, contract_name, date, freq, ZL_flag):
        hfbar_raw = self.connector.get_file(date, product_bardata_path + '/' + contract_name + '/' + date + '.pkl')
        datetimes = hfbar_raw.date.tolist()
        datetimes = [x + '000' for x in datetimes]
        hfbar_raw.index = pd.to_datetime(datetimes, format='%Y-%m-%d %H:%M:%S:%f').tolist()
        hfbar_raw_resampled = hfbar_raw.resample(freq, convention='end', label='right', closed='right'). \
            agg({
            'close': 'last',
            'low': 'min',
            'open': 'first',
            'high': 'max',
            'vol': 'sum',
            'syl2': 'last',
            'yclose': 'first',
            'syl1': 'last',
            'cjbs': 'sum',
            'amount': 'sum',
            'StockID': 'last',

        })
        hfbar_raw_resampled.replace(0, np.NaN, inplace=True)
        hfbar_raw_resampled.dropna(how='all', inplace=True)
        hfbar_raw_resampled.fillna(0, inplace=True)
        hfbar_raw_resampled['ZL'] = ZL_flag
        self.bar_Data = None

        return hfbar_raw_resampled

    def makeBarData(self, product_name, begt, endt, freq, hf_price, toRemote=False):
        """
        :param product_name: 品种名称
        :param freq: bar的频率
        :return:
        """
        product_bardata_local = info.DATAPATH.dataRootPath + '\\' + 'future_trade_data'
        product_bardata_remote = info.DATAPATH.server_root_path + '/' + 'future_trade_data'
        # ZL_info = pd.read_pickle(info.DATAPATH.dataRootPath + '\\' + 'future_basic_data' + '\\' + 'ZL.pkl')

        ZL_info = self.connector.get_file('ZL', info.DATAPATH.server_root_path + '/' + 'future_basic_data' + '/' + 'ZL.pkl')
        ZL_info = ZL_info[product_name].dropna()
        ZL_info = ZL_info.to_frame(name='ZL_name')
        ZL_info['last_ZL_name'] = ZL_info.ZL_name.shift(1)[ZL_info.ZL_name != ZL_info.ZL_name.shift(1)]
        ZL_info['last_ZL_name'][0] = np.NaN
        ZL_info.loc[ZL_info.last_ZL_name.isnull(), 'last_ZL_name'] = ZL_info.loc[ZL_info.last_ZL_name.isnull(), 'ZL_name']
        ZL_info = ZL_info.loc[begt:endt, :]

        if ('S' in freq) or ('min' in freq):
            product_bardata_path = product_bardata_remote + '/' + 'hfBar' + '/' + product_name

            resampled_results_list = []
            for index, row in ZL_info.iterrows():

                ZL_contract_name = row['ZL_name']
                last_ZL_contract_name = row['last_ZL_name']
                remote_filenames_index = self.connector.sftp.listdir(product_bardata_path + '/' + ZL_contract_name)
                remote_filenames_index = [x[:-4] for x in remote_filenames_index]
                if index in remote_filenames_index:
                    print(ZL_contract_name, index)
                    if ZL_contract_name==last_ZL_contract_name:
                        remote_filenames = self.connector.sftp.listdir(product_bardata_path + '/' + ZL_contract_name)
                        available_dates = [x[:-4] for x in remote_filenames]
                        yesterday = self.calendar.shift(1).loc[index, 'TradingDate']

                        hfbar_raw_resampled = self.read_hfData(
                            product_bardata_path=product_bardata_path,
                            contract_name=ZL_contract_name,
                            date=index,
                            freq=freq,
                            ZL_flag=True)

                        hfbar_raw_resampled = hfbar_raw_resampled.loc[:index + ' 16:00:00', :]
                        if yesterday in available_dates:
                            hfbar_raw_resampled_y = self.read_hfData(
                                product_bardata_path=product_bardata_path,
                                contract_name=ZL_contract_name,
                                date=yesterday,
                                freq=freq,
                                ZL_flag=True)
                            hfbar_raw_resampled_y = hfbar_raw_resampled_y.loc[yesterday + ' 21:00:00':, :]
                            hfbar_raw_resampled = pd.concat([hfbar_raw_resampled_y, hfbar_raw_resampled])
                            hfbar_raw_resampled['date'] = index

                        else:
                            pass
                        hfbar_raw_resampled = hfbar_raw_resampled[~hfbar_raw_resampled.index.duplicated(
                            keep='first'
                        )]
                        resampled_results_list.append(hfbar_raw_resampled)
                    elif ZL_contract_name!=last_ZL_contract_name:

                        remote_filenames = self.connector.sftp.listdir(product_bardata_path + '/' + ZL_contract_name)
                        available_dates = [x[:-4] for x in remote_filenames]
                        yesterday = self.calendar.shift(1).loc[index, 'TradingDate']


                        hfbar_raw_resampled_last_ZL = self.read_hfData(
                            product_bardata_path=product_bardata_path,
                            contract_name=last_ZL_contract_name,
                            date=index,
                            freq=freq,
                            ZL_flag=False)
                        hfbar_raw_resampled_last_ZL = hfbar_raw_resampled_last_ZL.loc[:index + ' 16:00:00', :]

                        hfbar_raw_resampled_ZL = self.read_hfData(
                            product_bardata_path=product_bardata_path,
                            contract_name=ZL_contract_name,
                            date=index,
                            freq=freq,
                            ZL_flag=True)
                        hfbar_raw_resampled_ZL = hfbar_raw_resampled_ZL.loc[:index + ' 16:00:00', :]

                        if yesterday in available_dates:
                            hfbar_raw_resampled_last_ZL_y = self.read_hfData(
                                product_bardata_path=product_bardata_path,
                                contract_name=last_ZL_contract_name,
                                date=yesterday,
                                freq=freq,
                                ZL_flag=False)
                            hfbar_raw_resampled_last_ZL_y = hfbar_raw_resampled_last_ZL_y.loc[yesterday + ' 21:00:00':, :]
                            hfbar_raw_resampled_last_ZL = pd.concat([hfbar_raw_resampled_last_ZL_y, hfbar_raw_resampled_last_ZL])
                            hfbar_raw_resampled_last_ZL['date'] = index
                            hfbar_raw_resampled_last_ZL = hfbar_raw_resampled_last_ZL[~hfbar_raw_resampled_last_ZL.index.duplicated(
                                keep='first'
                            )]

                            hfbar_raw_resampled_ZL_y = self.read_hfData(
                                product_bardata_path=product_bardata_path,
                                contract_name=ZL_contract_name,
                                date=yesterday,
                                freq=freq,
                                ZL_flag=True)
                            hfbar_raw_resampled_ZL_y = hfbar_raw_resampled_ZL_y.loc[yesterday + ' 21:00:00':, :]
                            hfbar_raw_resampled_ZL = pd.concat(
                                [hfbar_raw_resampled_ZL_y, hfbar_raw_resampled_ZL])
                            hfbar_raw_resampled_ZL['date'] = index
                            hfbar_raw_resampled_ZL = hfbar_raw_resampled_ZL[~hfbar_raw_resampled_ZL.index.duplicated(
                                keep='first'
                            )]
                        else:
                            pass
                        combined_hfbar_raw_resampled = pd.concat([hfbar_raw_resampled_last_ZL, hfbar_raw_resampled_ZL])
                        resampled_results_list.append(combined_hfbar_raw_resampled)

            hfbar_raw_resampled_total = pd.concat(resampled_results_list)
            hfbar_raw_resampled_total.index = pd.to_datetime(hfbar_raw_resampled_total.index.tolist()).tolist()
            hfbar_raw_resampled_total = hfbar_raw_resampled_total[~hfbar_raw_resampled_total.index.duplicated(keep='first')]
            hfbar_raw_resampled_total.to_pickle(product_bardata_local + '\\' + 'hfBar_ZL' + '\\' + product_name + '_' + freq + '.pkl')
            self.bar_Data = hfbar_raw_resampled_total
            if toRemote:
                local_path = 'tmp\\' + product_name + '_' + freq + '.pkl'
                hfbar_raw_resampled_total.to_pickle('tmp\\' + product_name + '_' + freq + '.pkl')
                remote_path = info.DATAPATH.server_root_path + '/' + 'future_trade_data' + '/'  + 'hfBar_ZL' + '/' + product_name + '_' + freq + '.pkl'
                self.connector.put_file(local_path, remote_path)
                os.remove(local_path)

        elif 'D' in freq:
            product_bardata_path_remote = product_bardata_remote + '/' + 'dailyBar' + '/' + product_name
            results_list = []
            for index, row in ZL_info.iterrows():
                ZL_contract_name = row['ZL_name']
                last_ZL_contract_name = row['last_ZL_name']
                if ZL_contract_name == last_ZL_contract_name:
                    contract_multiple = self.contract_info.loc[self.contract_info['代码'] == ZL_contract_name, '合约乘数'].values[0]
                    daily_raw = self.connector.get_file(ZL_contract_name, remote_path=product_bardata_path_remote + '/' + ZL_contract_name + '.pkl')
                    daily_raw.index = daily_raw.date.tolist()
                    daily_raw['ZL'] = True
                    if index in daily_raw.index.tolist():
                        daily_data = daily_raw.loc[[index], :]
                        if hf_price:
                            hd_raw = self.connector.get_file(index,
                                                             remote_path='/srv/data/future/future_trade_data/hfBar/' +
                                                                         product_name + '/' + ZL_contract_name + '/' + index + '.pkl')
                            hd_raw_clearing = hd_raw.iloc[-1199:, :]
                            twap = hd_raw.close.mean()
                            vwap = hd_raw.amount.sum() / hd_raw.vol.sum() / contract_multiple
                            twap_clearing = hd_raw_clearing.close.mean()
                            vwap_clearing = hd_raw_clearing.amount.sum() / hd_raw_clearing.vol.sum() / contract_multiple
                            daily_data.loc[index, 'twap'] = twap
                            daily_data.loc[index, 'vwap'] = vwap
                            daily_data.loc[index, 'twap_clearing'] = twap_clearing
                            daily_data.loc[
                                index, 'vwap_clearing'] = vwap_clearing if vwap_clearing is np.NaN else twap_clearing

                        results_list.append(daily_data)
                    else:
                        pass


                elif ZL_contract_name != last_ZL_contract_name:
                    daily_raw = self.connector.get_file(last_ZL_contract_name + '.pkl',
                                                        remote_path=product_bardata_path_remote + '/' + last_ZL_contract_name + '.pkl')
                    daily_raw.index = daily_raw.date.tolist()
                    daily_raw['ZL'] = False
                    results_list.append(daily_raw.loc[[index], :])
                    daily_raw = self.connector.get_file(ZL_contract_name + '.pkl',
                                                        remote_path=product_bardata_path_remote + '/' + ZL_contract_name + '.pkl')
                    daily_raw.index = daily_raw.date.tolist()
                    daily_raw['ZL'] = True
                    results_list.append(daily_raw.loc[[index], :])

            dailybar_raw_resampled_total = pd.concat(results_list)
            dailybar_raw_resampled_total.index = pd.to_datetime(dailybar_raw_resampled_total.index.tolist(), format='%Y-%m-%d').tolist()
            dailybar_raw_resampled_total.to_pickle(product_bardata_local + '\\' + 'dailyBar_ZL' + '\\' + product_name + '_' + freq + '.pkl')
            self.bar_Data = dailybar_raw_resampled_total
            print(dailybar_raw_resampled_total)
            if toRemote:
                local_path = 'tmp\\' + product_name + '_' + freq + '.pkl'
                dailybar_raw_resampled_total.to_pickle('tmp\\' + product_name + '_' + freq + '.pkl')
                remote_path = info.DATAPATH.server_root_path + '/' + 'future_trade_data' + '/'  + 'dailyBar_ZL' + '/' + product_name + '_' + freq + '.pkl'
                self.connector.put_file(local_path, remote_path)
                os.remove(local_path)

        def get_Bar_data():
            return self.bar_Data


if __name__ == '__main__':
    begt = '2010-01-01'
    endt = '2022-06-24'

    a = DataLoaderModule()
    df = pd.read_pickle(info.DATAPATH.dataRootPath + '\\' + 'future_basic_data' + '\\' + 'ProductNames.pkl')
    product_names = df['名称'].tolist()
    if '晚籼稻' in product_names:
        product_names.remove('晚籼稻')
    if '硬麦' in product_names:
        product_names.remove('硬麦')
    # product_names = ['镍']
    ni_index = product_names.index('镍')
    product_names = product_names[ni_index:]
    product_names = ['原油']
    freqs = ['10S', '15S', '30S', '1min', '5min', '10min', '15min']
    # freqs = ['1min']
    for product_name in product_names:
        for freq in freqs:
            a.makeBarData(product_name, begt, endt, freq=freq, hf_price=False, toRemote=True)

